//
// Created by ECNU on 2024/2/22.
//

#ifndef NANO_FAISS_INDEXFLAT_H
#define NANO_FAISS_INDEXFLAT_H

#include <nanofaiss/IndexFlatCodes.h>
#include <vector>

namespace faiss {

/** Index that stores full vectors and perform exhausted search */
struct IndexFlat : IndexFlatCodes {
    explicit IndexFlat(
            idx_t d,
            MetricType metric = MetricType::METRIC_L2);

    IndexFlat() = default;

    void search(
            idx_t n,
            const float *x,
            idx_t k,
            float *distances,
            idx_t *labels) const override;

    /// just copy x => bytes
    void sa_encode(idx_t n, const float *x, uint8_t *bytes) const override;

    /// just copy bytes => x
    void sa_decode(idx_t n, const uint8_t *bytes, float *x) const override;


    float* get_xb() {
        return reinterpret_cast<float*>(codes.data());
    }

    const float* get_xb() const {
        return reinterpret_cast<const float*>(codes.data());
    }
};

struct IndexFlatL2 : IndexFlat {

    /// if using the cached_l2norm, then inner product version l2 will be used
    /// |a - b|^2 = a^2 + b^2 - ab
    std::vector<float> cached_l2norm;

    explicit IndexFlatL2(idx_t d) : IndexFlat(d, MetricType::METRIC_L2) {}
    IndexFlatL2() {}

    // compute L2 norm
    void sync_l2norms();

    // clear L2 norm
    void clear_l2norms();


};

}

#endif // NANO_FAISS_INDEXFLAT_H
